#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MAE性能数据处理脚本
用于处理FuzzyTrust仿真生成的MAE数据，生成符合reference_rules.md要求的CSV文件

作者: FuzzyTrust项目组
日期: 2024
"""

import pandas as pd
import numpy as np
import os
import glob
from pathlib import Path

def process_mae_data(data_dir="examples/FuzzyVeins/results/collected_data/1"):
    """
    处理MAE数据并生成最终的性能CSV文件
    
    Args:
        data_dir: MAE数据文件所在目录
    """
    
    print(f"正在处理目录: {data_dir}")
    
    # 查找所有MAE性能数据文件
    performance_files = glob.glob(os.path.join(data_dir, "mae_performance_*.csv"))
    
    if not performance_files:
        print("未找到MAE性能数据文件")
        return
    
    print(f"找到 {len(performance_files)} 个性能数据文件")
    
    # 读取并合并所有性能数据
    all_data = []
    for file in performance_files:
        try:
            df = pd.read_csv(file)
            all_data.append(df)
            print(f"已读取文件: {file}, 数据行数: {len(df)}")
        except Exception as e:
            print(f"读取文件 {file} 时出错: {e}")
    
    if not all_data:
        print("没有成功读取任何数据文件")
        return
    
    # 合并所有数据
    combined_data = pd.concat(all_data, ignore_index=True)
    print(f"合并后总数据行数: {len(combined_data)}")
    
    # 按恶意节点比例分组并计算统计信息
    grouped = combined_data.groupby('malicious_ratio')
    
    final_results = []
    
    for ratio, group in grouped:
        print(f"\n处理恶意节点比例: {ratio}%")
        print(f"该比例下的数据行数: {len(group)}")
        
        # 计算各方法的MAE均值和标准差
        result = {
            'malicious_ratio': ratio,
            'mae_it2': group['mae_it2'].mean(),
            'std_it2': group['mae_it2'].std(),
            'mae_t1': group['mae_t1'].mean(),
            'std_t1': group['mae_t1'].std(),
            'mae_threshold': group['mae_threshold'].mean(),
            'std_threshold': group['mae_threshold'].std(),
            'mae_none': group['mae_none'].mean(),
            'std_none': group['mae_none'].std()
        }
        
        # 处理NaN值（当只有一个样本时标准差为NaN）
        for key in result:
            if pd.isna(result[key]):
                result[key] = 0.0
        
        final_results.append(result)
        
        print(f"  IT2-Sigmoid MAE: {result['mae_it2']:.4f} ± {result['std_it2']:.4f}")
        print(f"  Type-1 Sigmoid MAE: {result['mae_t1']:.4f} ± {result['std_t1']:.4f}")
        print(f"  Threshold MAE: {result['mae_threshold']:.4f} ± {result['std_threshold']:.4f}")
        print(f"  No Defense MAE: {result['mae_none']:.4f} ± {result['std_none']:.4f}")
    
    # 创建最终结果DataFrame
    final_df = pd.DataFrame(final_results)
    final_df = final_df.sort_values('malicious_ratio')
    
    # 保存最终结果
    output_file = os.path.join(data_dir, "mae_performance_final.csv")
    final_df.to_csv(output_file, index=False, float_format='%.6f')
    
    print(f"\n最终结果已保存到: {output_file}")
    print("\n最终结果预览:")
    print(final_df.to_string(index=False, float_format='%.4f'))
    
    return final_df

def process_raw_data(data_dir="examples/FuzzyVeins/results/collected_data/1"):
    """
    处理原始MAE数据并生成详细分析
    
    Args:
        data_dir: MAE数据文件所在目录
    """
    
    print(f"\n正在处理原始数据目录: {data_dir}")
    
    # 查找所有原始数据文件
    raw_files = glob.glob(os.path.join(data_dir, "mae_raw_data_*.csv"))
    
    if not raw_files:
        print("未找到原始MAE数据文件")
        return
    
    print(f"找到 {len(raw_files)} 个原始数据文件")
    
    # 读取并合并所有原始数据
    all_raw_data = []
    for file in raw_files:
        try:
            df = pd.read_csv(file)
            all_raw_data.append(df)
            print(f"已读取原始文件: {file}, 数据行数: {len(df)}")
        except Exception as e:
            print(f"读取原始文件 {file} 时出错: {e}")
    
    if not all_raw_data:
        print("没有成功读取任何原始数据文件")
        return
    
    # 合并所有原始数据
    combined_raw = pd.concat(all_raw_data, ignore_index=True)
    print(f"合并后原始数据总行数: {len(combined_raw)}")
    
    # 按恶意节点比例和评估方法分组分析
    analysis_results = []
    
    for ratio in sorted(combined_raw['malicious_ratio'].unique()):
        ratio_data = combined_raw[combined_raw['malicious_ratio'] == ratio]
        
        for method in ['IT2', 'Type1', 'Threshold', 'None']:
            method_data = ratio_data[ratio_data['evaluation_method'] == method]
            
            if len(method_data) > 0:
                analysis = {
                    'malicious_ratio': ratio,
                    'evaluation_method': method,
                    'sample_count': len(method_data),
                    'mae_mean': method_data['absolute_error'].mean(),
                    'mae_std': method_data['absolute_error'].std(),
                    'mae_min': method_data['absolute_error'].min(),
                    'mae_max': method_data['absolute_error'].max(),
                    'mae_median': method_data['absolute_error'].median(),
                    'trust_est_mean': method_data['estimated_trust'].mean(),
                    'trust_gt_mean': method_data['ground_truth_trust'].mean()
                }
                
                analysis_results.append(analysis)
    
    # 保存详细分析结果
    analysis_df = pd.DataFrame(analysis_results)
    analysis_output = os.path.join(data_dir, "mae_detailed_analysis.csv")
    analysis_df.to_csv(analysis_output, index=False, float_format='%.6f')
    
    print(f"\n详细分析结果已保存到: {analysis_output}")
    
    return analysis_df

def generate_summary_report(data_dir="examples/FuzzyVeins/results/collected_data/1"):
    """
    生成MAE性能评估总结报告
    
    Args:
        data_dir: MAE数据文件所在目录
    """
    
    print("\n=== MAE性能评估总结报告 ===")
    
    # 处理性能数据
    performance_df = process_mae_data(data_dir)
    
    if performance_df is not None and len(performance_df) > 0:
        print("\n1. 各方法在不同恶意节点比例下的MAE性能:")
        
        for _, row in performance_df.iterrows():
            ratio = row['malicious_ratio']
            print(f"\n恶意节点比例 {ratio}%:")
            print(f"  IT2-Sigmoid:    {row['mae_it2']:.4f} ± {row['std_it2']:.4f}")
            print(f"  Type-1 Sigmoid: {row['mae_t1']:.4f} ± {row['std_t1']:.4f}")
            print(f"  固定阈值法:      {row['mae_threshold']:.4f} ± {row['std_threshold']:.4f}")
            print(f"  无防御基线:      {row['mae_none']:.4f} ± {row['std_none']:.4f}")
        
        # 找出最佳性能
        print("\n2. 性能排名分析:")
        
        methods = ['mae_it2', 'mae_t1', 'mae_threshold', 'mae_none']
        method_names = ['IT2-Sigmoid', 'Type-1 Sigmoid', '固定阈值法', '无防御基线']
        
        for _, row in performance_df.iterrows():
            ratio = row['malicious_ratio']
            mae_values = [row[method] for method in methods]
            sorted_indices = np.argsort(mae_values)
            
            print(f"\n恶意节点比例 {ratio}% 性能排名（MAE越小越好）:")
            for i, idx in enumerate(sorted_indices):
                print(f"  {i+1}. {method_names[idx]}: {mae_values[idx]:.4f}")
    
    # 处理原始数据分析
    analysis_df = process_raw_data(data_dir)
    
    if analysis_df is not None and len(analysis_df) > 0:
        print("\n3. 数据统计信息:")
        total_samples = analysis_df['sample_count'].sum()
        unique_ratios = len(analysis_df['malicious_ratio'].unique())
        unique_methods = len(analysis_df['evaluation_method'].unique())
        
        print(f"  总样本数: {total_samples}")
        print(f"  恶意节点比例数: {unique_ratios}")
        print(f"  评估方法数: {unique_methods}")
    
    print("\n=== 报告生成完成 ===")

if __name__ == "__main__":
    import sys
    
    # 默认数据目录
    default_dir = "examples/FuzzyVeins/results/collected_data/1"
    
    # 如果提供了命令行参数，使用指定目录
    data_directory = sys.argv[1] if len(sys.argv) > 1 else default_dir
    
    print(f"MAE数据处理脚本启动")
    print(f"数据目录: {data_directory}")
    
    # 检查目录是否存在
    if not os.path.exists(data_directory):
        print(f"错误: 目录 {data_directory} 不存在")
        print("请确保已运行仿真并生成了MAE数据文件")
        sys.exit(1)
    
    # 生成完整的总结报告
    generate_summary_report(data_directory)
    
    print("\n处理完成！")
    print(f"请查看 {data_directory}/mae_performance_final.csv 文件")
    print("该文件可直接用于Origin绘制性能图")